#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
消融实验代码
验证不同特征组合对轴承故障诊断模型性能的影响
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import StratifiedKFold, cross_val_score
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.metrics import accuracy_score, f1_score, classification_report
import warnings
warnings.filterwarnings('ignore')

# 设置中文字体和图表样式
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False
plt.rcParams['figure.dpi'] = 150

class AblationStudy:
    def __init__(self, data_path=r'results/bearing_features.csv', n_splits=5, random_state=42):
        self.data_path = data_path
        self.n_splits = n_splits
        self.random_state = random_state
        self.results = {}
        
        # 定义特征组
        self.feature_groups = {
            'time_domain': ['rms', 'std', 'mean', 'skewness', 'kurtosis', 
                           'shape_factor', 'crest_factor', 'impulse_factor'],
            'freq_domain': ['freq_std', 'spectral_centroid', 'peak_frequency', 'freq_max',
                           'band_1_energy_ratio', 'band_3_energy_ratio', 'band_4_energy_ratio'],
            'envelope': ['envelope_mean', 'envelope_std', 'envelope_max'],
            'fault_freq': ['BPFO_energy', 'BPFI_energy', 'BSF_energy', 'FR_energy'],
            'advanced': ['sample_entropy', 'approximate_entropy', 'fractal_dimension',
                        'teager_kaiser_energy', 'wavelet_energy']
        }
        
        # 加载数据
        self.load_data()
        
    def load_data(self):
        """加载和预处理数据"""
        print("=== 加载数据 ===")
        df = pd.read_csv(self.data_path)
        
        # 数据清洗
        df = df.fillna(0).replace([np.inf, -np.inf], np.nan).fillna(0)
        
        # 编码标签
        le = LabelEncoder()
        df['label_encoded'] = le.fit_transform(df['label'])
        self.class_names = le.classes_
        
        # 获取所有可用特征 - 排除非数值列
        exclude_cols = ['label', 'filename', 'rpm', 'sampling_rate', 'label_encoded', 
                       'bearing_type', 'signal_length']
        self.all_features = [col for col in df.columns if col not in exclude_cols]
        
        # 过滤存在的特征并确保是数值类型
        self.available_features = []
        for f in self.all_features:
            if f in df.columns:
                # 检查列是否为数值类型
                if pd.api.types.is_numeric_dtype(df[f]):
                    self.available_features.append(f)
                else:
                    print(f"跳过非数值特征: {f}")
        
        # 确保所有特征都是数值类型
        self.X = df[self.available_features].astype(float)
        self.y = df['label_encoded']
        
        print(f"数据形状: {self.X.shape}")
        print(f"类别: {self.class_names}")
        print(f"可用特征数: {len(self.available_features)}")
        
        # 检查是否有无效值
        if self.X.isnull().any().any():
            print("警告: 发现NaN值，将用0填充")
            self.X = self.X.fillna(0)
        
        if np.isinf(self.X).any().any():
            print("警告: 发现无穷值，将用0填充")
            self.X = self.X.replace([np.inf, -np.inf], 0)
        
    def get_feature_combinations(self):
        """定义消融实验的特征组合"""
        combinations = {}
        
        # 1. 基线：所有特征
        combinations['all_features'] = self.available_features
        
        # 2. 单组特征
        for group_name, features in self.feature_groups.items():
            available_group_features = [f for f in features if f in self.available_features]
            if available_group_features:
                combinations[f'only_{group_name}'] = available_group_features
        
        # 3. 两两组合
        group_names = list(self.feature_groups.keys())
        for i in range(len(group_names)):
            for j in range(i+1, len(group_names)):
                group1, group2 = group_names[i], group_names[j]
                features1 = [f for f in self.feature_groups[group1] if f in self.available_features]
                features2 = [f for f in self.feature_groups[group2] if f in self.available_features]
                if features1 and features2:
                    combinations[f'{group1}_{group2}'] = features1 + features2
        
        # 4. 逐步添加特征组
        combinations['time_freq'] = []
        for group in ['time_domain', 'freq_domain']:
            features = [f for f in self.feature_groups[group] if f in self.available_features]
            combinations['time_freq'].extend(features)
        
        combinations['time_freq_envelope'] = combinations['time_freq'].copy()
        envelope_features = [f for f in self.feature_groups['envelope'] if f in self.available_features]
        combinations['time_freq_envelope'].extend(envelope_features)
        
        combinations['time_freq_envelope_fault'] = combinations['time_freq_envelope'].copy()
        fault_features = [f for f in self.feature_groups['fault_freq'] if f in self.available_features]
        combinations['time_freq_envelope_fault'].extend(fault_features)
        
        # 5. 基于重要性的特征选择
        # 使用RandomForest计算特征重要性
        rf = RandomForestClassifier(n_estimators=100, random_state=self.random_state)
        rf.fit(self.X, self.y)
        feature_importance = rf.feature_importances_
        
        # 按重要性排序
        importance_df = pd.DataFrame({
            'feature': self.available_features,
            'importance': feature_importance
        }).sort_values('importance', ascending=False)
        
        # 选择不同数量的重要特征
        for n in [5, 10, 15, 20]:
            if n <= len(importance_df):
                top_features = importance_df.head(n)['feature'].tolist()
                combinations[f'top_{n}_features'] = top_features
        
        # 6. 基于特征类型的组合（新增）
        # 时域+频域组合
        time_freq_features = []
        for group in ['time_domain', 'freq_domain']:
            features = [f for f in self.feature_groups[group] if f in self.available_features]
            time_freq_features.extend(features)
        if time_freq_features:
            combinations['time_domain_freq_domain'] = time_freq_features
        
        # 频域+故障频率组合
        freq_fault_features = []
        for group in ['freq_domain', 'fault_freq']:
            features = [f for f in self.feature_groups[group] if f in self.available_features]
            freq_fault_features.extend(features)
        if freq_fault_features:
            combinations['freq_domain_fault_freq'] = freq_fault_features
        
        return combinations
    
    def evaluate_feature_combination(self, features, combination_name):
        """评估特定特征组合的性能"""
        print(f"评估特征组合: {combination_name} ({len(features)} 个特征)")
        
        # 检查特征是否都存在
        missing_features = [f for f in features if f not in self.available_features]
        if missing_features:
            print(f"  警告: 缺失特征 {missing_features}")
            features = [f for f in features if f in self.available_features]
        
        if not features:
            print(f"  错误: 没有可用特征")
            return None
        
        # 准备数据
        X_subset = self.X[features]
        
        # 标准化
        scaler = StandardScaler()
        X_scaled = scaler.fit_transform(X_subset)
        
        # 交叉验证
        rf = RandomForestClassifier(
            n_estimators=100,
            max_depth=10,
            random_state=self.random_state,
            class_weight='balanced'
        )
        
        # 计算多个指标
        cv_scores = {}
        
        # 定义评估指标
        metrics = {
            'accuracy': 'accuracy',
            'f1': 'f1_weighted',
            'precision': 'precision_weighted',
            'recall': 'recall_weighted'
        }
        
        # 计算各种指标
        for metric_name, metric_func in metrics.items():
            try:
                scores = cross_val_score(rf, X_scaled, self.y, cv=self.n_splits, 
                                       scoring=metric_func, n_jobs=-1)
                cv_scores[metric_name] = {
                    'mean': np.mean(scores),
                    'std': np.std(scores),
                    'scores': scores
                }
            except Exception as e:
                print(f"  警告: 计算{metric_name}时出错: {e}")
                cv_scores[metric_name] = {
                    'mean': 0.0,
                    'std': 0.0,
                    'scores': [0.0] * self.n_splits
                }
        
        result = {
            'combination_name': combination_name,
            'features': features,
            'n_features': len(features),
            'cv_scores': cv_scores
        }
        
        print(f"  准确率: {cv_scores['accuracy']['mean']:.4f} ± {cv_scores['accuracy']['std']:.4f}")
        print(f"  F1分数: {cv_scores['f1']['mean']:.4f} ± {cv_scores['f1']['std']:.4f}")
        
        return result
    
    def run_ablation_study(self):
        """运行完整的消融实验"""
        print("=== 开始消融实验 ===")
        
        # 获取特征组合
        combinations = self.get_feature_combinations()
        print(f"定义了 {len(combinations)} 个特征组合")
        
        # 评估每个组合
        results = []
        for name, features in combinations.items():
            result = self.evaluate_feature_combination(features, name)
            if result:
                results.append(result)
        
        self.results = results
        return results
    
    def plot_ablation_results(self, save_path='ablation_results.png'):
        """绘制消融实验结果"""
        if not self.results:
            print("没有结果可绘制")
            return
        
        # 准备数据
        df_results = []
        for result in self.results:
            for metric in ['accuracy', 'f1', 'precision', 'recall']:
                df_results.append({
                    'combination': result['combination_name'],
                    'n_features': result['n_features'],
                    'metric': metric,
                    'mean': result['cv_scores'][metric]['mean'],
                    'std': result['cv_scores'][metric]['std']
                })
        
        df_plot = pd.DataFrame(df_results)
        
        # 创建图表
        fig, axes = plt.subplots(2, 2, figsize=(20, 15))
        fig.suptitle('消融实验结果分析', fontsize=16)
        
        # 1. 不同特征组合的性能对比
        ax1 = axes[0, 0]
        accuracy_data = df_plot[df_plot['metric'] == 'accuracy'].sort_values('mean', ascending=False)
        bars = ax1.bar(range(len(accuracy_data)), accuracy_data['mean'], 
                      yerr=accuracy_data['std'], capsize=5, alpha=0.7)
        ax1.set_xticks(range(len(accuracy_data)))
        ax1.set_xticklabels(accuracy_data['combination'], rotation=45, ha='right')
        ax1.set_ylabel('准确率')
        ax1.set_title('不同特征组合的准确率对比')
        ax1.grid(True, alpha=0.3)
        
        # 添加数值标签
        for i, (bar, mean, std) in enumerate(zip(bars, accuracy_data['mean'], accuracy_data['std'])):
            ax1.text(bar.get_x() + bar.get_width()/2, bar.get_height() + std + 0.01,
                    f'{mean:.3f}', ha='center', va='bottom', fontsize=8)
        
        # 2. 特征数量 vs 性能
        ax2 = axes[0, 1]
        for metric in ['accuracy', 'f1']:
            metric_data = df_plot[df_plot['metric'] == metric]
            ax2.errorbar(metric_data['n_features'], metric_data['mean'], 
                        yerr=metric_data['std'], marker='o', label=metric, capsize=5)
        ax2.set_xlabel('特征数量')
        ax2.set_ylabel('性能分数')
        ax2.set_title('特征数量 vs 模型性能')
        ax2.legend()
        ax2.grid(True, alpha=0.3)
        
        # 3. 不同指标的热力图
        ax3 = axes[1, 0]
        pivot_data = df_plot.pivot(index='combination', columns='metric', values='mean')
        sns.heatmap(pivot_data, annot=True, fmt='.3f', cmap='YlOrRd', ax=ax3)
        ax3.set_title('不同组合的指标热力图')
        ax3.set_xlabel('评估指标')
        ax3.set_ylabel('特征组合')
        
        # 4. 特征组贡献分析
        ax4 = axes[1, 1]
        group_performance = {}
        for result in self.results:
            name = result['combination_name']
            if 'only_' in name:
                group = name.replace('only_', '')
                group_performance[group] = result['cv_scores']['accuracy']['mean']
        
        if group_performance:
            groups = list(group_performance.keys())
            scores = list(group_performance.values())
            bars = ax4.bar(groups, scores, alpha=0.7, color=['skyblue', 'lightgreen', 'lightcoral', 'gold'])
            ax4.set_ylabel('准确率')
            ax4.set_title('各特征组的单独贡献')
            ax4.grid(True, alpha=0.3)
            
            # 添加数值标签
            for bar, score in zip(bars, scores):
                ax4.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.01,
                        f'{score:.3f}', ha='center', va='bottom')
        
        plt.tight_layout()
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.show()
        
        return fig
    
    def generate_report(self, save_path='ablation_report.txt'):
        """生成消融实验报告"""
        if not self.results:
            print("没有结果可报告")
            return
        
        with open(save_path, 'w', encoding='utf-8') as f:
            f.write("=== 轴承故障诊断消融实验报告 ===\n\n")
            
            # 实验设置
            f.write("1. 实验设置\n")
            f.write(f"   数据集: {self.data_path}\n")
            f.write(f"   交叉验证折数: {self.n_splits}\n")
            f.write(f"   特征组合数: {len(self.results)}\n")
            f.write(f"   总特征数: {len(self.available_features)}\n\n")
            
            # 结果排序
            f.write("2. 特征组合性能排名\n")
            f.write("   按准确率排序:\n")
            sorted_results = sorted(self.results, key=lambda x: x['cv_scores']['accuracy']['mean'], reverse=True)
            
            for i, result in enumerate(sorted_results, 1):
                f.write(f"   {i:2d}. {result['combination_name']:25s} "
                       f"({result['n_features']:2d} 特征) "
                       f"准确率: {result['cv_scores']['accuracy']['mean']:.4f} ± {result['cv_scores']['accuracy']['std']:.4f}\n")
            
            # 最佳组合详情
            f.write("\n3. 最佳特征组合详情\n")
            best_result = sorted_results[0]
            f.write(f"   组合名称: {best_result['combination_name']}\n")
            f.write(f"   特征数量: {best_result['n_features']}\n")
            f.write(f"   特征列表: {', '.join(best_result['features'])}\n")
            f.write(f"   准确率: {best_result['cv_scores']['accuracy']['mean']:.4f} ± {best_result['cv_scores']['accuracy']['std']:.4f}\n")
            f.write(f"   F1分数: {best_result['cv_scores']['f1']['mean']:.4f} ± {best_result['cv_scores']['f1']['std']:.4f}\n")
            f.write(f"   精确率: {best_result['cv_scores']['precision']['mean']:.4f} ± {best_result['cv_scores']['precision']['std']:.4f}\n")
            f.write(f"   召回率: {best_result['cv_scores']['recall']['mean']:.4f} ± {best_result['cv_scores']['recall']['std']:.4f}\n")
            
            # 特征组分析
            f.write("\n4. 特征组贡献分析\n")
            group_scores = {}
            for result in self.results:
                if 'only_' in result['combination_name']:
                    group = result['combination_name'].replace('only_', '')
                    group_scores[group] = result['cv_scores']['accuracy']['mean']
            
            if group_scores:
                sorted_groups = sorted(group_scores.items(), key=lambda x: x[1], reverse=True)
                for group, score in sorted_groups:
                    f.write(f"   {group:15s}: {score:.4f}\n")
            
            # 结论和建议
            f.write("\n5. 结论和建议\n")
            f.write("   基于消融实验结果，可以得出以下结论:\n")
            f.write("   - 最佳特征组合为: {}\n".format(best_result['combination_name']))
            f.write("   - 最优特征数量为: {}\n".format(best_result['n_features']))
            f.write("   - 建议在后续模型中使用此特征组合\n")
            f.write("   - 可以考虑进一步优化特征选择策略\n")
        
        print(f"消融实验报告已保存到: {save_path}")
    
    def save_results_csv(self, save_path='ablation_results.csv'):
        """保存结果到CSV文件"""
        if not self.results:
            print("没有结果可保存")
            return
        
        # 准备数据
        data = []
        for result in self.results:
            row = {
                'combination_name': result['combination_name'],
                'n_features': result['n_features'],
                'features': ';'.join(result['features'])
            }
            
            # 添加各种指标
            for metric in ['accuracy', 'f1', 'precision', 'recall']:
                row[f'{metric}_mean'] = result['cv_scores'][metric]['mean']
                row[f'{metric}_std'] = result['cv_scores'][metric]['std']
            
            data.append(row)
        
        df_results = pd.DataFrame(data)
        df_results = df_results.sort_values('accuracy_mean', ascending=False)
        df_results.to_csv(save_path, index=False, encoding='utf-8-sig')
        print(f"消融实验结果已保存到: {save_path}")

def main():
    """主函数"""
    print("=== 轴承故障诊断消融实验 ===")
    
    # 创建消融实验对象
    ablation = AblationStudy()
    
    # 运行消融实验
    results = ablation.run_ablation_study()
    
    # 绘制结果
    ablation.plot_ablation_results()
    
    # 生成报告
    ablation.generate_report()
    
    # 保存CSV结果
    ablation.save_results_csv()
    
    print("\n=== 消融实验完成 ===")
    print("生成的文件:")
    print("- ablation_results.png: 可视化结果")
    print("- ablation_report.txt: 详细报告")
    print("- ablation_results.csv: 数据结果")
    
    return results

if __name__ == "__main__":
    main()
